import json
import os
import sys
from pathlib import Path

import cv2
import numpy as np

from camera.camera import Camera, CameraBaseTransforamtion

# from model.transforms import denormalize_end, normalize_joint
from robots.piper import PiperArm


class DynaCali:
    """
    Dynamically adjust the transition matrix by manual clicking
    """

    def __init__(self, camera: Camera, arm: PiperArm):
        self.image = None
        self.original_image = None
        self.window_name = "Pixel calibration"
        self.scale = 1.0
        self.offset_x = 0
        self.offset_y = 0
        self.dragging = False
        self.last_x, self.last_y = 0, 0
        self.pixel = None
        self.actual_end = None
        self.camera_xyz = []
        self.gripper_xyz = []
        self.colors = [
            (0, 0, 255),  # 红色
            (0, 255, 0),  # 绿色
            (255, 0, 0),  # 蓝色
            (0, 255, 255),  # 黄色
            (255, 0, 255),  # 紫色
            (255, 255, 0),  # 青色
        ]

        self.arm = arm  # xyz-joint mapping should be well trained!
        self.camera = camera
        self.transform = CameraBaseTransforamtion(f"camera/{self.camera.model}/{self.arm.can}", suffix="")
        self.lr = 0.05

    def capture_camera_image(self):
        return self.camera.capture_current_frame()

    def camera_pixel_to_3d(self, x, y):
        return self.camera.get_3d_point_from_pixel(x, y)

    def load_image(self):
        """加载图像"""

        self.original_image = self.capture_camera_image()
        if self.original_image is None:
            print(f"Error: No image captured!")
            return False
        self.image = self.original_image.copy()
        self.scale = 1.0
        self.offset_x = 0
        self.offset_y = 0
        self.pixels = []
        return True

    def display_coordinates(self, event, x, y, flags, param):
        """鼠标事件回调函数"""
        # 将显示坐标转换为原始图像坐标
        original_x = int((x + self.offset_x) / self.scale)
        original_y = int((y + self.offset_y) / self.scale)

        # 确保坐标在有效范围内
        if (
            original_x < 0
            or original_x >= self.original_image.shape[1]
            or original_y < 0
            or original_y >= self.original_image.shape[0]
        ):
            return

        if event == cv2.EVENT_LBUTTONDOWN:
            # left-click: get the true end pos
            self.pixel = [original_x, original_y]
            self.c_xyz = self.camera_pixel_to_3d(original_x, original_y)
            self.est_end = self.transform(self.c_xyz)

            # self.arm.go_zero()
            # raise to make it safer
            middle_xyz = self.arm.get_effector_xeuler()[:3]
            middle_xyz[2] += 0.05
            self.arm.move_to(middle_xyz)
            self.arm.move_to(self.est_end)
            self.actual_end = self.arm.get_effector_xeuler()[:3]

            print(f"================================================================")
            print(f"Camera_xyz={self.c_xyz}")
            print(f"Try to move to_xyz={self.est_end}")
            print(f"Actually move to_xyz={self.actual_end}")

            actual_joint = self.arm.get_joint_status()
            print(f"Actually move to joits={actual_joint}")
            print(f"================================================================")
            # if self.actual_end is None:
            #     # self.arm.go_zero()
            #     self.arm.raise_and_move_to(self.est_end)
            #     self.actual_end = self.arm.get_end_pos()[:3]
            #     print(f"Camera_xyz={self.c_xyz}")
            #     print(f"Try to move to_xyz={self.est_end}")
            #     print(f"Actually move to_xyz={self.actual_end}")
            # else:

        elif event == cv2.EVENT_RBUTTONDOWN:
            self.actual_end = None
            self.pixel = None

        elif event == cv2.EVENT_MBUTTONDOWN:
            # 中键点击：开始拖动
            self.dragging = True
            self.last_x, self.last_y = x, y

        elif event == cv2.EVENT_MOUSEMOVE:
            # 鼠标移动
            if self.dragging:
                # 拖动图像
                dx = x - self.last_x
                dy = y - self.last_y
                self.offset_x += dx
                self.offset_y += dy
                self.last_x, self.last_y = x, y

        elif event == cv2.EVENT_MBUTTONUP:
            # 中键释放：结束拖动
            self.dragging = False

        # 显示当前鼠标位置的坐标
        cv2.displayOverlay(self.window_name, f"Pixel: ({original_x}, {original_y})", 100)

    def run(self):
        """运行标定工具"""

        # image_path = self.camera.save_current_image()
        if not self.load_image():
            return
        img_path, img = self.camera.save_current_image()
        cv2.imread(img_path)
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.window_name, self.display_coordinates)

        while True:
            # 创建当前视图的副本
            display = self.capture_camera_image()

            if self.pixel:
                px, py = self.pixel
                cv2.circle(display, (px, py), 2, (0, 0, 255), -1)
                cv2.putText(
                    display, f"Target=({px}, {py})", (px + 10, py - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1
                )

            # if self.revised_pixel:
            #     px, py = self.revised_pixel
            #     cv2.circle(display, (px, py), 2, (255, 0, 0), -1)
            #     cv2.putText(display, f"Modified=({px}, {py})",
            #                 (px-10, py+10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)

            # 应用缩放和偏移
            h, w = display.shape[:2]
            view_x1 = max(0, int(self.offset_x / self.scale))
            view_y1 = max(0, int(self.offset_y / self.scale))
            view_x2 = min(w, int((self.offset_x + w) / self.scale))
            view_y2 = min(h, int((self.offset_y + h) / self.scale))

            view = display[view_y1:view_y2, view_x1:view_x2]
            view = cv2.resize(view, (w, h))

            # 显示图像
            cv2.imshow(self.window_name, view)

            # 处理键盘事件
            key = cv2.waitKey(1) & 0xFF

            if key == 27:  # ESC键退出
                break
            elif key == ord("+") or key == ord("="):  # 放大
                self.scale *= 1.1
            elif key == ord("-"):  # 缩小
                self.scale /= 1.1
            elif key == ord("r"):
                # Gradient descent
                # if self.actual_end is not None:
                #     print(f"Perform gradient descent !")
                #     T = self.transform.camera_base_matrix
                #     c_hom = np.append(self.c_xyz, 1)[: ,None]
                #     gap = (self.est_end - self.actual_end)[None, :]

                #     self.transform.camera_base_matrix = T - self.lr * c_hom @ gap
                #     # update transformation matrix and set the act_end as the tar
                #     self.est_end = self.transform(self.c_xyz)
                #     # self.arm.go_zero()
                #     self.arm.raise_and_move_to(self.est_end)
                #     self.actual_end = self.arm.get_end_pos()[:3]
                #     print(f"Camera_xyz={self.c_xyz}")
                #     print(f"Try to move to_xyz={self.est_end}")
                #     print(f"Actually move to_xyz={self.actual_end}")

                self.scale = 1.0
                self.offset_x = 0
                self.offset_y = 0
            elif key == ord("s"):  # Save the matrix
                self.save_matrix()
        cv2.destroyAllWindows()

    def save_matrix(self):

        f_dir = f"data/calibrations/{self.arm.can}"
        if not os.path.exists(f_dir):
            os.mkdir(f_dir)
        f_name = os.path.join(f_dir, "cali_matrix_dyna.json")

        with open(f_name, "w") as f:
            json.dump(self.transform.camera_base_matrix.tolist(), f)

        print(f"Saved transformation matrix to {f_name}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Use: python pixel_calibration.py <can>")
        sys.exit(1)

    arm = PiperArm(sys.argv[1], wait_execution=True)
    calibrator = DynaCali(Camera(), arm)
    calibrator.run()


# python calibration_dyna.py can0
